// cuda/conv2d_optim_kernel_fp32.cu
#include <torch/extension.h>
#include "conv2d_fp32.h" // 定义param_t和DTYPE

// 前向传播核函数（直接卷积方法）
__global__ void implgemm(param_t param)
{
    // 计算线程在输出特征图扁平后的索引（Oh * Ow方向）
    int x_out = blockIdx.x * blockDim.x + threadIdx.x; 
    // 计算线程在输出通道（K）方向的索引
    int y_out = blockIdx.y * blockDim.y + threadIdx.y; 
    // 计算线程在批次（N）方向的索引
    int z_out = blockIdx.z;                            

    // 检查线程索引是否越界
    if (x_out >= param.Oh * param.Ow || y_out >= param.k || z_out >= param.n)
    {
        return;
    }

    // 从扁平索引中分解出输出特征图的行和列索引
    int oh = x_out / param.Ow;
    int ow = x_out % param.Ow;

    // 计算输入特征图块的左上角坐标（考虑步长和填充）
    int h_in_start = oh * param.u - param.p;
    int w_in_start = ow * param.v - param.q;

    DTYPE current_sum = 0.0f;

    // 遍历输入通道
    for (int c_in = 0; c_in < param.c; ++c_in) {        
        // 遍历卷积核行（滤波器高度R）
        for (int kr = 0; kr < param.r; ++kr) {        
            // 遍历卷积核列（滤波器宽度S）
            for (int ks = 0; ks < param.s; ++ks) {    
                
                // 计算当前输入特征图的坐标（考虑偏移和填充）
                int h_in = h_in_start + kr;
                int w_in = w_in_start + ks;

                // 边界检查，确保输入坐标有效
                if (h_in >= 0 && h_in < param.h && w_in >= 0 && w_in < param.w) {
                    // 计算输入张量索引：[批次, 输入通道, 高, 宽]
                    int input_idx = z_out * param.c * param.h * param.w +   // 批次偏移
                                    c_in  * param.h * param.w +         // 输入通道偏移
                                    h_in  * param.w +                   // 行偏移
                                    w_in;                               // 列偏移

                    // 计算权重张量索引：[输出通道, 输入通道, 核高, 核宽]
                    int weight_idx = y_out * param.c * param.r * param.s +  // 输出通道（滤波器k）偏移
                                     c_in  * param.r * param.s +        // 输入通道（滤波器c）偏移
                                     kr    * param.s +                  // 卷积核行偏移
                                     ks;                                // 卷积核列偏移
                    
                    // 累加乘积累积结果
                    current_sum += param.input[input_idx] * param.weight[weight_idx];
                }
            }
        }
    }

    // 计算输出张量索引：[批次, 输出通道, 输出高, 输出宽]
    int output_idx = z_out * param.k * param.Oh * param.Ow +  // 批次偏移
                     y_out * param.Oh * param.Ow +            // 输出通道K偏移
                     oh    * param.Ow +                       // 输出行Oh偏移
                     ow;                                      // 输出列Ow偏移
    // 写入输出结果
    param.output[output_idx] = current_sum;
}

void conv2d_cuda_forward(param_t param)
{
    // 定义线程块维度（简化处理，类似基线直接卷积）
    // 每个线程块在输出特征图扁平维度上处理的元素数（最大32）
    int threadx = min(32, (param.Oh * param.Ow + 31) / 32);  
    // 每个线程块在输出通道维度上处理的通道数（最大为CHANNEL_TILE）
    int thready = min(CHANNEL_TILE, param.k);                 
    int threadz = 1; // 批次维度线程数设为1                                           
    
    // 计算网格维度（自动向上取整）
    int grid_dim_x = (param.Oh * param.Ow + threadx - 1) / threadx;
    int grid_dim_y = (param.k + thready - 1) / thready;
    int grid_dim_z = param.n; // 每个批次单独处理
    
    // 计算每个SM（流式多处理器）能容纳的最大活跃线程块数
    int max_active_blocks_per_sm;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks_per_sm, 
        implgemm_optim, // 优化后的核函数（注：原代码中函数名可能有误，应为implgemm）
        threadx * thready * threadz, // 每个线程块的线程数
        0 // 共享内存大小（此处未使用）
    );
    
    // 获取设备属性以计算SM数量
    cudaDeviceProp props;
    cudaGetDeviceProperties(&props, 0);
    int sm_count = props.multiProcessorCount;
    int target_blocks = sm_count * max_active_blocks_per_sm; // 目标线程块总数
    
    // 调整网格大小以优化SM利用率
    if (grid_dim_x * grid_dim_y * grid_dim_z < target_blocks) {
        // 计算扩展因子以接近目标线程块数
        int scale_factor = (target_blocks + (grid_dim_x * grid_dim_y * grid_dim_z - 1)) / 
                          (grid_dim_x * grid_dim_y * grid_dim_z);
        grid_dim_x = min(grid_dim_x * scale_factor, (param.Oh * param.Ow + threadx - 1) / threadx);
    }
    
    dim3 block(threadx, thready, threadz);
    dim3 grid(grid_dim_x, grid_dim_y, grid_dim_z);
    
    // 启动优化后的卷积核（注：原函数名implgemm_optim可能应为implgemm）
    implgemm_optim<<<grid, block>>>(param);
    
    // 检查CUDA调用错误
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        printf("CUDA错误: %s\n", cudaGetErrorString(err));
    }
}


// 反向传播数据核函数（计算输入梯度）
__global__ void implgemmbwddata(param_t param)
{
    uint32_t tx = threadIdx.x; // 线程在块内的x维度索引
    int bx = blockIdx.x;       // 块在x维度的索引
    int by = blockIdx.y;       // 块在y维度的索引

    const uint32_t lane_id = threadIdx.x % 32;    // 线程在warp中的lane索引（0-31）
    const uint32_t warp_id = threadIdx.x / 32;    // warp索引
    const uint32_t mma_tid_x = (lane_id / 2) % 8; // 矩阵乘法操作的x维度线程索引
    const uint32_t mma_tid_y = (lane_id / 16) * 2 + (lane_id % 2); // y维度线程索引

    // 共享内存地址计算（权重和梯度输出）
    uint32_t weight_lds_addr = (warp_id / 2) * 32 + mma_tid_y * 4;
    uint32_t gradoutput_lds_addr = (warp_id % 2) * 64 + mma_tid_x * 4;

    // 计算输入特征图索引（用于梯度计算）
    int x = bx * 128 + gradoutput_lds_addr;
    int y = by * 128 + weight_lds_addr;
    int z = blockIdx.z; // 批次索引

    // 共享内存声明（存储梯度输出和权重）
    __shared__ DTYPE smemgradoutput[8 * 128];
    __shared__ DTYPE smemweight[8 * 132];

    DTYPE weight_ldg_reg[4];  // 寄存器变量存储加载的权重
    DTYPE gradoutput_ldg_reg[4]; // 寄存器变量存储加载的梯度输出

    int posOh_ori[4]; // 原始输出特征图行坐标（用于边界计算）
    int posOw_ori[4]; // 原始输出特征图列坐标（用于边界计算）
#pragma unroll // 展开循环以提高效率
    for (int i = 0; i < 4; ++i)
    {
        // 计算输出特征图坐标（考虑步长和填充反向传播）
        posOh_ori[i] = ((bx * 128 + tx % 32 + i * 32) / param.w) - (param.r - 1 - param.p);
        posOw_ori[i] = ((bx * 128 + tx % 32 + i * 32) % param.w) - (param.s - 1 - param.q);
    }

    // 偏移量计算（输出/权重张量）
    int outOffset = z * param.k * param.Oh * param.Ow; // 批次偏移
    int weiC = (by * 128 + tx / 8 * 4); // 当前处理的输入通道索引
    int outKOffset = param.Oh * param.Ow; // 输出通道内的偏移
    int weiCOffset = param.r * param.s; // 输入通道内的权重偏移
    int weiKOffset = param.c * param.r * param.s; // 输出通道间的权重偏移

    // 共享内存存储地址计算
    uint32_t weight_sts_addr = (tx % 8) * 132 + (tx / 8) * 4; // 权重存储地址
    uint32_t gradoutput_sts_addr = (tx / 32) * 128 + (tx % 32); // 梯度输出存储地址

    DTYPE weight_frag[8];       // 权重片段（用于矩阵乘法）
    DTYPE gradoutput_frag[8];   // 梯度输出片段（用于矩阵乘法）
    DTYPE gradinput_frag[8][8]; // 输入梯度片段（累加结果）
#pragma unroll
    for (int i = 0; i < 8; ++i)
    {
#pragma unroll
        for (int j = 0; j < 8; ++j)
        {
            gradinput_frag[i][j] = 0.0f; // 初始化累加器
        }
    }

    // 遍历所有卷积核元素和输出通道
    for (int krs = 0; krs < param.r * param.s * param.k; krs += 8)
    {
        int curKRS = krs + tx % 8; // 当前处理的核元素索引
        int rs = param.r * param.s - 1 - curKRS % (param.r * param.s); // 反转核索引（对应反向传播）
        int curK = curKRS / (param.r * param.s); // 当前处理的输出通道索引

#pragma unroll
        for (int i = 0; i < 4; ++i)
        {
            // 边界检查并加载权重（处理越界情况）
            if ((curK * param.r * param.s + rs) < param.r * param.s * param.k)
            {
                weight_ldg_reg[i] = param.weight[curK * weiKOffset + (weiC + i) * weiCOffset + rs];
            }
            else
            {
                weight_ldg_reg[i] = 0.0f; // 越界时置零
            }
        }

        int curK2 = (krs + tx / 32) / (param.r * param.s); // 当前处理的输出通道（按warp分组）
        int curR = ((krs + tx / 32) % (param.r * param.s)) / param.s; // 当前核行索引
        int curS = ((krs + tx / 32) % (param.r * param.s)) % param.s; // 当前核列索引

#pragma unroll
        for (int i = 0; i < 4; ++i)
        {
            // 计算输出特征图坐标并检查边界
            int curOh = posOh_ori[i] + curR; 
            int curOw = posOw_ori[i] + curS; 
            int outOffsetTmp = curK2 * outKOffset + curOh * param.Ow + curOw;
            if (curOh >= 0 && curOw >= 0 && curOw < param.Ow && curOh < param.Oh)
            {
                gradoutput_ldg_reg[i] = param.grad_output[outOffset + outOffsetTmp]; // 加载梯度输出
            }
            else
            {
                gradoutput_ldg_reg[i] = 0.0f; // 越界时置零
            }
        }

        // 将数据存入共享内存（同步前需确保所有线程完成写入）
        for (int i = 0; i < 4; ++i)
        {
            smemweight[weight_sts_addr + i] = weight_ldg_reg[i];
        }
        for (int i = 0; i < 4; ++i)
        {
            smemgradoutput[gradoutput_sts_addr + i * 32] = gradoutput_ldg_reg[i];
        }
        __syncthreads(); // 等待所有线程完成共享内存写入

#pragma unroll
        for (int subkrs = 0; subkrs < 8; ++subkrs)
        {
#pragma unroll
            for (int i = 0; i < 4; ++i)
            {
                // 从共享内存加载权重和梯度输出片段
                weight_frag[i] = smemweight[weight_lds_addr + subkrs * 132 + i];
                weight_frag[i + 4] = smemweight[weight_lds_addr + subkrs * 132 + i + 16];
            }
#pragma unroll
            for (int i = 0; i < 4; ++i)
            {
                gradoutput_frag[i] = smemgradoutput[gradoutput_lds_addr + subkrs * 128 + i];
                gradoutput_frag[i + 4] = smemgradoutput[gradoutput_lds_addr + subkrs * 128 + i + 32];
            }

            // 矩阵乘法累加（计算输入梯度）
#pragma unroll
            for (int i = 0; i < 8; ++i)
            {
#pragma unroll
                for (int j = 0; j < 8; ++j)
                {
                    gradinput_frag[i][j] += (weight_frag[i] * gradoutput_frag[j]);
                }
            }
        }
        __syncthreads(); // 等待所有线程完成当前迭代
    }

    // 将累加结果写入输入梯度张量
    int gradinputOffset;
#pragma unroll
    for (int i = 0; i < 4; ++i)
    {
#pragma unroll
        for (int j = 0; j < 4; ++j)
        {
            // 计算输入梯度索引（考虑批次、通道、空间坐标）
            gradinputOffset = z * param.c * param.h * param.w + (y + i) * param.h * param.w + x + j;
            if (x + j < param.h * param.w && y + i < param.c)
            {
                param.grad_input[gradinputOffset] = gradinput_frag[i][j]; // 写入低32位部分
            }
            gradinputOffset = z * param.c * param.h * param.w + (y + i) * param.h * param.w + x + j + 32;
            if (x + j + 32 < param.h * param.w && y + i < param.c)
            {
                param.grad_input[gradinputOffset] = gradinput_frag[i][j + 4]; // 写入高32位部分
            }
            gradinputOffset = z * param.c * param.h * param.w + (y + i + 16) * param.h * param.w + x + j;
            if (x + j < param.h * param.w && y + i + 16 < param.c)
            {
                param.grad_input[gradinputOffset] = gradinput_frag[i + 4][j]; // 写入下16通道部分
            }
            gradinputOffset = z * param.c * param.h * param.w + (y + i + 16) * param.h * param.w + x + j + 32;
            if (x + j + 32 < param.h * param.w && y + i + 16 < param.c)
            {
                param.grad_input[gradinputOffset] = gradinput_frag[i + 4][j + 4]; // 写入综合部分
            }
        }
    }
}

// 反向传播权重核函数（计算权重梯度）
__global__ void implgemmbwdweight(param_t param)
{
    uint32_t tx = threadIdx.x; // 线程在块内的x维度索引
    int bx =